import numpy as np
import torch
from torch import nn

from config import Config


def robust(_data):
    switch = int(np.random.rand() * 2) + 1
    phase = int(np.random.rand() * Config["chunk"])
    phase -= Config["chunk"]
    phase = phase // switch
    if phase >= 0:
        data = np.append(_data[phase:], _data[:phase]*0)
    else:
        data = np.append(_data[0:-phase]*0, _data[:phase])
    return data


class Encode(object):
    def __init__(self, robust_flag=False):
        super(Encode, self).__init__()
        self.robust_flag = robust_flag

    def __call__(self, wav_data, target=None):
        wav_data = np.frombuffer(wav_data, dtype=np.int16)
        wav_data = np.float32(wav_data)
        mm = max(abs(wav_data))+1
        if mm > 1000.0:
            wav_data = wav_data / mm
            if target and self.robust_flag:
                wav_data = robust(wav_data)
        else:
            wav_data *= 0.0
        target = np.float32(target)
        wav_data = torch.from_numpy(np.array([wav_data]))
        target = torch.from_numpy(np.array([target]))
        return wav_data, target


class Decode(object):
    def __init__(self):
        super(Decode, self).__init__()

    def __call__(self, inputs, targets):
        loss = []
        assert (targets is not None)
        assert inputs.shape == targets.shape
        # loss = nn.L1Loss()(inputs, targets)
        for i in range(len(inputs)):
            loss_ = nn.L1Loss()(inputs[i], targets[i])
            loss.append(loss_)

        return loss
